import pandas as pd
import json
import requests
import time
import os
from tqdm import tqdm


# 1. 加载本地TSV文件
def load_darkbench_tsv(file_path):
    """解析DarkBench的TSV文件"""
    try:
        # 尝试不同的分隔符格式
        df = pd.read_csv(file_path, sep='\t', header=None, names=['category', 'prompt'], skiprows=1)
        print(f"成功加载 {len(df)} 条Prompt，示例类别: {df['category'][0]}")
        return df
    except Exception as e:
        print(f"TSV文件解析失败: {str(e)}")
        print("请检查文件格式是否正确（应为制表符分隔的TSV）")
        return None


# DeepSeek API 调用函数
def get_deepseek_response(prompt, api_key, model="deepseek-chat", max_retries=3):
    """通过API获取DeepSeek的响应"""
    url = "https://api.deepseek.com/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0,
        "max_tokens": 1024
    }

    for attempt in range(max_retries):
        try:
            response = requests.post(url, headers=headers, json=payload, timeout=30)
            response.raise_for_status()  # 检查HTTP错误

            # 检查响应是否有效
            if "choices" in response.json() and len(response.json()["choices"]) > 0:
                return response.json()["choices"][0]["message"]["content"]
            else:
                print(f"响应格式异常: {response.text}")
                return None

        except requests.exceptions.Timeout:
            print(f"请求超时（尝试 {attempt + 1}/{max_retries}）")
            time.sleep(5)  # 等待后重试
        except Exception as e:
            print(f"API调用失败: {str(e)}")
            return None

    print(f"请求失败，已达到最大重试次数 ({max_retries})")
    return None


# 保存进度函数
def save_progress(responses, filename="responses.json"):
    """实时保存进度"""
    try:
        with open(filename, "w", encoding="utf-8") as f:
            json.dump(responses, f, ensure_ascii=False, indent=2)
        return True
    except Exception as e:
        print(f"保存失败: {str(e)}")
        return False


if __name__ == "__main__":
    # 加载数据
    dataset = load_darkbench_tsv("darkbench.tsv")

    if dataset is None:
        exit(1)

    # 你的DeepSeek API密钥
    API_KEY = "sk-208664f290e74a0dae34d10536b51cff"  # 替换成你的真实key

    # 检查是否已有部分结果
    responses = []
    if os.path.exists("responses.json"):
        try:
            with open("responses.json", "r", encoding="utf-8") as f:
                responses = json.load(f)
            print(f"检测到已有 {len(responses)} 条响应，将继续追加")
        except:
            print("无法读取现有响应文件，将重新开始")

    # 获取已处理ID列表
    processed_ids = {item['id'] for item in responses}

    # 设置进度条
    total_items = min(660, len(dataset))  # 最多处理50条
    progress_bar = tqdm(total=total_items, desc="处理进度")

    # 更新进度条到当前进度
    progress_bar.update(len(responses))

    try:
        # 遍历数据集
        for i, row in dataset.iterrows():
            # 只处理前50条
            if i >= total_items:
                break

            # 跳过已处理的
            if i in processed_ids:
                progress_bar.update(1)
                continue

            try:
                # 调用API获取响应
                response_text = get_deepseek_response(row['prompt'], API_KEY)

                if response_text is None:
                    print(f"第{i}条响应获取失败，跳过")
                    responses.append({
                        "id": i,
                        "category": row['category'],
                        "prompt": row['prompt'],
                        "response": "ERROR: FAILED TO GET RESPONSE"
                    })
                else:
                    responses.append({
                        "id": i,
                        "category": row['category'],
                        "prompt": row['prompt'],
                        "response": response_text
                    })

                # 实时保存
                if save_progress(responses):
                    print(f"已保存第{i}条结果")

                # 更新进度条
                progress_bar.update(1)

                # 添加延迟避免速率限制
                time.sleep(1.5)  # 1.5秒间隔

            except Exception as e:
                print(f"处理第{i}条失败: {str(e)}")
                responses.append({
                    "id": i,
                    "category": row['category'],
                    "prompt": row['prompt'],
                    "response": f"ERROR: {str(e)}"
                })
                save_progress(responses)
                progress_bar.update(1)

    finally:
        # 确保最后保存
        save_progress(responses)
        progress_bar.close()
        print(f"处理完成！共保存 {len(responses)} 条响应")